import baostock as bs
import pandas as pd

class PerformanceTest:
    def __init__(self, initial_codes, start_date, end_date):
        self.holdings = {}
        self.start_date = start_date
        self.start_year = int(start_date[:4])
        self.end_year = int(end_date[:4])
        self.end_date = end_date
        self.total_value = 0.0
        self.trade_log = {}

        # 获取start_date的股价
        code_dict = {}
        open_sum = 0.0
        for code in initial_codes:
            rs = bs.query_history_k_data_plus(code, "code,open", start_date=self.start_date,
                                              end_date=self.end_date, frequency="d", adjustflag="3")
            data_list = []
            while (rs.error_code == '0') & rs.next():
                # 获取一条记录，将记录合并在一起
                data_list.append(rs.get_row_data())
            result = pd.DataFrame(data_list, columns=rs.fields)
            init_price = float(result[code]['k_data']['open'][0])
            code['initial_price'] = init_price
            open_sum += init_price
            code_dict[code] = {
                'init_price' : init_price,
            }
        # 建仓
        for code in initial_codes:
            self.holdings[code] = {
                'amount': open_sum / code_dict[code]['initial_price'] * 100,
                'code': code,
                'cost_price': code_dict[code]['initial_price']
            }
            self.total_value += self.holdings[code]['amount'] * code_dict[code]['cost_price']

    def test(self):
        for year in range(self.start_year + 1, self.end_year + 1):
            # 检查pr
            for holding in self.holdings:

